import pymysql
from towhee import ops
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility

staticUrl = "https://s.shanhu.club"
config = {'towhee': {'model': 'resnet50', 'device': None}, 'milvus': {'host': '127.0.0.1', 'port': '19530',
    'dim': 2048, 'table_name': 'image_search', 'index_type': 'IVF_FLAT', 'metric_type': 'L2'},
    'mysql': {'host': "127.0.0.1", 'user': "root", 'password': "root", 'database': "ziyouxia", 'charset': "utf8"}
}


def create_milvus_collection(table_name, config):
    is_exist = utility.has_collection(table_name)
    fields = [
        FieldSchema(name='id', dtype=DataType.INT64, description='image id', is_primary=True, auto_id=False),
        FieldSchema(name='uniq_no', dtype=DataType.VARCHAR, description='image uniq_no', max_length=63),
        FieldSchema(name='name', dtype=DataType.VARCHAR, description='image name', max_length=256),
        FieldSchema(name='vec', dtype=DataType.FLOAT_VECTOR, description='image embedding vectors', dim=config['towhee']['dim'])
    ]
    schema = CollectionSchema(fields=fields, description='reverse image search')
    collection_table = Collection(name=table_name, schema=schema)
    if not is_exist:
        index_params = {
            'metric_type': config['milvus']['metric_type'],
            'index_type': config['milvus']['index_type'],
            'params': {"nlist": config['milvus']['dim']}
        }
        collection_table.create_index(field_name='vector', index_params=index_params)
    return collection_table


def get_list_from_mysql(config, offset, limit):
    conn = pymysql.connect(host=config["host"], user=config["user"], password=config["password"], database=config["database"], charset=config["charset"])
    cursor = conn.cursor()

    sql = f"select a.id,a.uniq_no,a.name,b.thumb from sp_reference as a left join uc_upfile as b on a.thumb=b.uniq_no order by id ASC limit {offset}, {limit}"
    list_items = []
    try:
        # 执行 sql 语句
        cursor.execute(sql)
        result = cursor.fetchall()
        for row in result:
            list_items.append({'id': row[0], 'uniq_no': row[1], 'name': row[2], 'ads': row[3]})
    except:
        print("Error: unable to fetch data")
    # 关闭数据库连接
    conn.close()
    return list_items


connections.connect(host=config['milvus']['HOST'], port=config['milvus']['PORT'])
dataTable = create_milvus_collection(config['milvus']['table_name'], config)

# 执行获取数据逻辑
offset = 0
limit = 100
while True:
    items = get_list_from_mysql(config["mysql"], offset, limit)
    if len(items) < 1:
        break
    item_id, uniq, name, item_vec = [], [], [], []
    for item in items:
        img = ops.image_decode()(staticUrl + item['ads'])
        vec = ops.image_embedding.timm(model_name=config['towhee']['model'], device=config['towhee']['device'])(img)
        item_id.append(item['id'])
        name.append(item['name'])
        uniq.append(item['uniq_no'])
        item_vec.append(vec)
        dataTable.insert([item_id, uniq, name, item_vec])
    if len(items) < limit:
        break

dataTable.flush()
dataTable.load()
print(dataTable)
# Check collection
print('Number of data inserted:', dataTable.num_entities)
